package org.piggy.core.algorithm;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.SortedMap;
import java.util.TreeMap;

public class ConsistentHash {

    private final int replicaCount;

    // 虚拟节点标识符前缀
    private final String prefix = "Virtual-";

    // 内部节点存储结构
    private final TreeMap<Long, String> ring = new TreeMap<>();

    public ConsistentHash() {
        this.replicaCount = 100;
    }

    public ConsistentHash(int replicaCount) {
        this.replicaCount = replicaCount;
    }

    /**
     * 添加物理节点
     * @param node 节点标识
     */
    public void addNode(String node) {
        // 生成虚拟节点并添加到哈希环
        for (int i = 0; i < replicaCount; i++) {
            String virtualNode = prefix + node + "-" + i;
            long hash = hash(virtualNode);
            ring.put(hash, node);
        }
    }

    /**
     * 移除节点（物理节点及其所有虚拟节点）
     * @param node 节点标识
     */
    public void removeNode(String node) {
        for (int i = 0; i < replicaCount; i++) {
            String virtualNode = prefix + node + "-" + i;
            long hash = hash(virtualNode);
            ring.remove(hash);
        }
    }

    /**
     * 获取数据对应的节点
     * @param dataKey 数据标识
     * @return 物理节点标识
     */
    public String getNode(String dataKey) {
        if (ring.isEmpty())
            return null;
        long dataHash = hash(dataKey);
        SortedMap<Long, String> tailMap = ring.tailMap(dataHash);
        long nodeHash = tailMap.isEmpty() ? ring.firstKey() : tailMap.firstKey();
        return ring.get(nodeHash);
    }

    // MD5哈希算法实现（128位截断为64位）
    private long hash(String input) {
        try {
            MessageDigest md = MessageDigest.getInstance("MD5");
            byte[] digest = md.digest(input.getBytes(StandardCharsets.UTF_8));
            long h = 0;
            // 取前64位作为long型哈希值
            for (int i = 0; i < 8; i++) {
                h <<= 8;
                h |= (digest[i] & 0xFF);
            }
            return h;
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("MD5算法不可用", e);
        }
    }

    // 测试用例
    public static void main(String[] args) {
        ConsistentHash ch = new ConsistentHash();
        // 添加物理节点（每个节点100个虚拟节点）
        ch.addNode("192.168.12.1:8080");
        ch.addNode("192.168.12.1:8085");
        ch.addNode("192.168.12.1:8090");

        // 模拟数据分配
        String[] testKeys = {"user:001", "order:007", "product:009"};
        for (String key : testKeys) {
            System.out.println("Key " + key + " 分配到节点: " + ch.getNode(key));
        }

        // 动态移除节点测试
        ch.removeNode("192.168.12.1:8085");
        System.out.println("\n移除后，重新分配:");
        for (String key : testKeys) {
            System.out.println("Key " + key + " 分配到节点: " + ch.getNode(key));
        }
    }
}
